#-*- coding: utf-8 -*-
#
# Copyright 2013 Antoine Drouin (poinix@gmail.com)
#
# This file is part of PAT.
#
#    PAT is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    PAT is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with PAT.  If not, see <http://www.gnu.org/licenses/>.
#

"""
This is the generic interface for an autopilot
"""
import math
import numpy as np
import xml.etree.ElementTree as ET
import StringIO
import importlib

def load(type, variant, _dm=None, _args=None, _time=None):
    mod_str = "pat.vehicles.{:s}.autopilot_{:s}".format(type, variant)
    ap = importlib.import_module(mod_str)
    return ap.Autopilot(_dm=_dm, _args=_args, _time=_time)


class Autopilot:
    """ An autopilot"""
    def __init__(self, _args=None, _dm=None, _time=None):
#        self.Ui = np.zeros(dm.iv_size)
#        self.refi = np.zeros(dm.sv_size)
        self.modes= []
        self.mode = 0
        self.ref = None

    def run(self, time, X):
        self.Ui, self.refi= self.modes[self.mode].run(time, X, self.ref)
        return self.Ui, self.refi
        
    def add_mode(self, mode):
        self.modes.append(mode)

    def get_idx_selected_mode(self): return self.mode

    def set_cur_mode(self, no):
        if len(self.modes[no].reference) > 0: self.ref = self.modes[no].reference[0]
        self.mode = no

    def get_mode(self, no): return self.modes[no]

    def get_cur_mode(self): return self.modes[self.mode]

    def set_param_ref(self, param, t0): self.ref.set_param(param, t0)

    def properties(self):
        return self.get_cur_mode().properties()

    def get_U_str(self):
        return "{:s}".format(self.Ui)

    def set_ref(self, ref): self.ref = ref


class Reference():
    """ A Reference"""
    def __init__(self, Xr, Ur, name="None"):
        self.trim = Xr, Ur
        self.name = name
        
    def get(self, time):
        return self.trim

class StepReference(Reference):
    def __init__(self, Xr, Ur, c=0, name=None, period=16., ampl=1.):
        Reference.__init__(self, Xr, Ur, name)
        self.period = period; self.ampl = ampl
        self.c = c # the coordinate of the sate vector on which the step is applied
        
    def get(self, time):
        Xr, Ur = self.trim
        Xr1, Ur1 = np.array(Xr), np.array(Ur)
        a = math.fmod(time, self.period)
        if a >= self.period/2. : Xr1[self.c] += self.ampl
        else: Xr1[self.c] -= self.ampl
        return Xr1, Ur1 
        

class ControlLaw:
    """ A Control Law"""
    def __init__(self, name="None", refs=()):
        self.name = name
        self.reference = refs

#    def run(self, time, X, ref=None):
#        return U, Xref, Uref

    def properties(self):
        return "None"


class LinearFeedback(ControlLaw):
    
    def __init__(self, sv_size, iv_size, Xr, Ur, K=None, name=None, refs=()):
        ControlLaw.__init__(self, name, refs)
        if K<> None: self.K = K
        else: self.K = np.zeros((iv_size, sv_size))
        self.Xr = np.array(Xr)
        self.Ur = np.array(Ur)
        
    def run(self, time, X, ref=None):
        if ref:
            self.Xr, self.Ur = ref.get(time)
        dX = X - self.Xr
        dU = np.dot(self.K, dX)
        self.U = self.Ur + dU
        return np.array(self.U), np.array(self.Xr)

    def set_gain(self, K): self.K = K

    def properties(self):
        return "K=\n{:s}".format(self.K)

    def load_from_file(self, filename):
        """XML parser"""
        print "ap loading xml from ",filename 
        tree = ET.parse(filename)
        root = tree.getroot()
#        self.name = root.attrib["name"]
        modes = root.find("mode")
        K = np.genfromtxt(StringIO.StringIO(modes.find("gain").text), delimiter=",")
        print K
        self.K = K
        

class LQRFeedback(LinearFeedback):
#    def __init__(self, sv_size, iv_size, Xr, Ur, K=None, name=None):
#        ap.LinearFeedback.__init__(self, dm.sv_size, dm.iv_size, Xr, Ur, K, name)

    def properties(self):
        np.set_printoptions(precision=3, suppress=True)
        props = "Q=diag({:s})\nR=diag({:s})\nK=\n{:s}".format(np.diag(self.Q), np.diag(self.R), self.K)
        return props

